/**
 * \file re.c
 **/

#include <stdio.h>
#include <string.h>
#include "re.h"
#include "mem.h"
#include "token.h"
#include "packrat.h"
#include "log.h"

/**
 * \brief Types of tokens used when compiling a regular expression.
 **/
#define RE_DOT 1       /* Token '.' */
#define RE_LETTER 2    /* A letter (void* is to a uni_char) */
#define RE_BOX 3       /* Token [...] (match many characters) */
#define RE_NOT_BOX 4   /* Token [^...] (match any except in []) */
#define RE_STAR 10     /* Token '*' */
#define RE_QUESTION 11 /* Token '?' */
#define RE_PLUS 12     /* Token '+' */
#define RE_PIPE 13     /* Token '|' */
#define RE_CONCAT 14   /* No token, but used in AST. */
#define RE_LPAREN 21   /* Token '(' */
#define RE_RPAREN 22   /* Token ')' */

/**
 * \brief AST node for a regular expression. If type = RE_DOT, no other fields
 * are used. If type = RE_LETTER, extra is a uni_char*. If type = RE_STAR,
 * RE_QUESTION or RE_PLUS, node1 is used. If type = RE_PIPE or RE_CONCAT, both
 * node1 and node2 are used.
 **/
struct RRENodeT{
  uint8_t type;
  union {
    void* extra;
    struct RRENodeT* node1;
  };
  struct RRENodeT* node2;
};
typedef struct RRENodeT RRENode;

/**
 * \brief An element of a [...] part of a regular expression. May be a match
 * of a single character, or a match of a range of characters a-b.
 **/
#define RE_SINGLE 1
#define RE_RANGE  2
struct RREBoxElementT{
  uint8_t type;      /* May be RE_SINGLE or RE_RANGE. */
  uni_char char1;
  uni_char char2;    /* Used if type = RE_RANGE. */
};
typedef struct RREBoxElementT RREBoxElement;

/**
 * \brief Nonterminals used in parsing a regular expression, and associated
 * matchers.
 **/
#define RE_EXPR   0    /* expr :=  expr ('+'|'?'|'*')  precedence: 3
                        *        | expr expr           precedence: 2
                        *        | expr '|' expr       precedence: 1 */
static RSemantic match_expr(RParsing* p, int precedence);


/**
 * \brief Construct a new RE AST node.
 **/
static inline RRENode* r_re_new_node(uint8_t type);

/**
 * \brief Investigate a RE AST tree.
 **/
static void r_re_ast_display(RRENode* root);

/**
 * \brief Recursive implementation of r_re_ast_display().
 **/
static void r_re_ast_display_r(RRENode* root, int depth);

/**
 * \brief Construct a ReSubExpression out of a parsed box. If inverted = true,
 * then match [^...]
 **/
static ReSubExpression r_re_box(RVector* box, bool inverted);

/**
 * \brief Tokenize and parse the [...] regular expressions.
 * \return true if regular expression well formed. Otherwise false.
 **/
static bool r_re_tokenize_box(RVector* token_list,
                               uint8_t** cur,
                               uint8_t* end);

/**
 * \brief Construct the NFA from the AST.
 **/
static ReSubExpression r_re_construct(RRENode* ast);

RSemantic match_expr(RParsing* p, int precedence)
{
  void *extra;
  RSemantic e1, e2;
  RState state;
  RRENode* ast;

  /* precedence: 3 */
  if (precedence <= 3){
    state = r_prat_save(p);
    if ((e1 = r_prat_match(p, RE_EXPR, 3))){
      if (r_prat_consume(p, RE_STAR, &extra)){
        ast = r_re_new_node(RE_STAR);
        ast->node1 = e1;
        return ast;
      }

      if (r_prat_consume(p, RE_PLUS, &extra)){
        ast = r_re_new_node(RE_PLUS);
        ast->node1 = e1;
        return ast;
      }
    }
    r_prat_restore(p, state);
  }

  /* precedence: 2 */
  if (precedence <= 2){
    state = r_prat_save(p);
    if ((e1 = r_prat_match(p, RE_EXPR, 2))){
      if ((e2 = r_prat_match(p, RE_EXPR, 2))){
        ast = r_re_new_node(RE_CONCAT);
        ast->node1 = e1;
        ast->node2 = e2;
        return ast;
      }
    }
    r_prat_restore(p, state);
  }

  /* precedence: 1 */
  if (precedence <= 1){
    state = r_prat_save(p);
    if ((e1 = r_prat_match(p, RE_EXPR, 1))){
      if (r_prat_consume(p, RE_PIPE, &extra)){
        if ((e2 = r_prat_match(p, RE_EXPR, 1))){
          ast = r_re_new_node(RE_PIPE);
          ast->node1 = e1;
          ast->node2 = e2;
          return ast;
        }
      }
    }
    r_prat_restore(p, state);
  }

  /* Match empty parentheses */
  state = r_prat_save(p);
  if (r_prat_consume(p, RE_LPAREN, &extra)){
    if (r_prat_consume(p, RE_RPAREN, &extra)){
      /* Empty parens. Consume and keep going. */
    } else {
      r_prat_restore(p, state);
    }
  }

  /* expr := '(' expr ')' */
  state = r_prat_save(p);
  if (r_prat_consume(p, RE_LPAREN, &extra)){
    if ((e1 = r_prat_match(p, RE_EXPR, 0))){
      if (r_prat_consume(p, RE_RPAREN, &extra)){
        return e1;
      }
      r_prat_restore(p, state);
    }
  }

  /* TODO: There must be a better way to do this: */
  if (r_prat_consume(p, RE_LETTER, &extra)){
    ast = r_re_new_node(RE_LETTER);
    ast->extra = extra;
    return ast;
  }
  if (r_prat_consume(p, RE_DOT, &extra)){
    ast = r_re_new_node(RE_DOT);
    return ast;
  }
  if (r_prat_consume(p, RE_BOX, &extra)){
    ast = r_re_new_node(RE_BOX);
    ast->extra = extra;
    return ast;
  }
  if (r_prat_consume(p, RE_NOT_BOX, &extra)){
    ast = r_re_new_node(RE_NOT_BOX);
    ast->extra = extra;
    return ast;
  }

  return NULL;
}

ReSubExpression r_re_alter(ReSubExpression first, ReSubExpression second)
{
  ReSubExpression result;
  NFANode* split_node = (NFANode*) r_mem_malloc(sizeof(NFANode));
  r_nfanode_init_split(split_node, 2);
  result.input = split_node;
  r_nfanode_add_child(split_node, first.input);
  r_nfanode_add_child(split_node, second.input);
  /* Now we want to connect the outputs of first and second into
   * result.output: */
  /* Find end of first linked list */
  NFANode* node = first.output;
  while (node->extra != NULL) node = node->extra;
  /* Append the outputs of second subexpression */
  node->extra = second.output;
  result.output = first.output;
  return result;
}

void r_re_ast_display(RRENode* root)
{
  r_re_ast_display_r(root, 0);
}

void r_re_ast_display_r(RRENode* root, int depth)
{
  int i;
  for (i = 0; i < depth; i++) printf(" ");

  switch(root->type){
    case RE_STAR:
      printf("*\n");
      r_re_ast_display_r(root->node1, depth+2);
      break;
    case RE_PLUS:
      printf("+\n");
      r_re_ast_display_r(root->node1, depth+2);
      break;
    case RE_CONCAT:
      printf("concat\n");
      r_re_ast_display_r(root->node1, depth+2);
      r_re_ast_display_r(root->node2, depth+2);
      break;
    case RE_PIPE:
      printf("|\n");
      r_re_ast_display_r(root->node1, depth+2);
      r_re_ast_display_r(root->node2, depth+2);
      break;
    case RE_LETTER:
      printf("%d\n", *((uni_char*) root->extra));
  }
}

NFANode* r_re_compile(char* expression, int term_id)
{
  /* Tokenize the expression */
  uint8_t* cur = (uint8_t*) expression;
  uint8_t* end = cur + strlen(expression);
  RVector token_list;
  r_tok_list_init(&token_list);

  while (cur != end){
    uni_char letter;
    uni_char* token_letter;

    if (r_utf8_read_uni_char(&cur, end, &letter) != R_SUCCESS)
      return NULL;

    switch(letter){
      case '.':
        r_tok_list_add(&token_list, RE_DOT, NULL);
        break;
      case '+':
        r_tok_list_add(&token_list, RE_PLUS, NULL);
        break;
      case '?':
        r_tok_list_add(&token_list, RE_QUESTION, NULL);
        break;
      case '*':
        r_tok_list_add(&token_list, RE_STAR, NULL);
        break;
      case '|':
        r_tok_list_add(&token_list, RE_PIPE, NULL);
        break;
      case '(':
        r_tok_list_add(&token_list, RE_LPAREN, NULL);
        break;
      case ')':
        r_tok_list_add(&token_list, RE_RPAREN, NULL);
        break;
      case '[':
        if (not r_re_tokenize_box(&token_list, &cur, end))
          return NULL;
        break;
      case '\\':
        /* Match the following letter: */
        if (r_utf8_read_uni_char(&cur, end, &letter) != R_SUCCESS)
          return NULL;
        /* proceed to default */
      default:
        /* Any letter. */
        token_letter = r_mem_alloca(sizeof(uni_char));
        *token_letter = letter;
        r_tok_list_add(&token_list, RE_LETTER, token_letter);
    }
  }

  /* Parse the expression. */
  RGrammar* grammar;
  grammar = r_prat_grammar_new();
  r_prat_grammar_add(grammar, RE_EXPR, match_expr);
  RParsing parsing;
  parsing = r_prat_new_parsing(grammar, token_list);
  RRENode* ast = r_prat_parse(&parsing, RE_EXPR);
  if (ast == NULL)
    return NULL;

  ReSubExpression sub_expr = r_re_construct(ast);
  NFANode* terminator = r_mem_malloc(sizeof(NFANode));
  r_nfanode_init_term(terminator, term_id);
  r_re_route(sub_expr.output, terminator);
  return sub_expr.input;
}

bool r_re_tokenize_box(RVector* token_list, uint8_t** cur, uint8_t* end)
{
  RVector* box = r_mem_malloc(sizeof(RVector));
  box->el_size = sizeof(RREBoxElement);
  r_vector_init(box);
  RREBoxElement element;

  /* Used for temporary read-aheads. */
  uint8_t* temp;
  uni_char next_letter;

  /* Check if first character is '^' */
  bool inverted = false;
  temp = *cur;
  if (r_utf8_read_uni_char(&temp, end, &next_letter) != R_SUCCESS)
    return false;
  if (next_letter == '^'){
    inverted = true;
    *cur = temp;
  }

  while (*cur != end){
    uni_char letter;

    if (r_utf8_read_uni_char(cur, end, &letter) != R_SUCCESS)
      return false;
    r_debug(R_DETAILED, "read letter %d", letter);

    if (letter == ']')
      break;

    if (letter == '\\'){
      if (r_utf8_read_uni_char(cur, end, &letter) != R_SUCCESS)
        return false;
    } /* Proceed further, having replaced letter. */

    /* Read next letter. */
    temp = *cur; /* Use this variable, because we don't want to
                  * change cur */
    int rv = r_utf8_read_uni_char(&temp, end, &next_letter);
    if (rv == R_SUCCESS){
      if (next_letter == '-'){
        /* Read one more letter. */
        if (r_utf8_read_uni_char(&temp, end, &next_letter) != R_SUCCESS)
          return false;

        /* Construct a range of characters. */
        element.type = RE_RANGE;
        element.char1 = letter;
        element.char2 = next_letter;
        r_vector_append(box, &element);
        *cur = temp;
        continue;
      }
    } else if (rv == R_UTF8_LIMITED){
      /* There is no next character. Ignore. */
    } else {
      /* An error occurred! */
      return false;
    }
    /* Not a range of characters. */
    element.type = RE_SINGLE;
    element.char1 = letter;
    r_vector_append(box, &element);
  } /* while (cur != end) */

  if (box->size == 0)
    return false; /* Having [] or [^] is invalid in the regular expression. */
  if (inverted)
    r_tok_list_add(token_list, RE_NOT_BOX, box);
  else
    r_tok_list_add(token_list, RE_BOX, box);
  return true;
}

ReSubExpression r_re_box(RVector* box, bool inverted)
{
  ReSubExpression result;
  NFANode* node = (NFANode*) r_mem_malloc(sizeof(NFANode));
  r_nfanode_init_char(node, box->size);
  if (inverted)
    node->type = R_NFA_NODE_NOT_CHAR;

  int i;
  r_debug(R_DETAILED, "constructing SubExpression from %d elements in box.",
           box->size);
  for (i = 0; i < box->size; i++){
    RREBoxElement box_element;
    r_vector_get_into(box, i, &box_element);
    switch (box_element.type){
      case RE_SINGLE:
        r_nfanode_add_class(node, R_NFA_CLASS_SINGLE, box_element.char1, 0);
        break;
      case RE_RANGE:
        r_nfanode_add_class(node, R_NFA_CLASS_RANGE, box_element.char1,
                             box_element.char2);
        break;
      default:
        assert_never();
    }
  }
  result.input = node;
  result.output = node;
  return result;
}

ReSubExpression r_re_construct(RRENode* ast)
{
  switch(ast->type){
    case RE_DOT:
      return r_re_dot();
    case RE_LETTER:
      return r_re_letter(*((uni_char*) ast->extra));
    case RE_QUESTION:
      return r_re_question(r_re_construct(ast->node1));
    case RE_PLUS:
      return r_re_plus(r_re_construct(ast->node1));
    case RE_STAR:
      return r_re_star(r_re_construct(ast->node1));
    case RE_CONCAT:
      return r_re_concat(r_re_construct(ast->node1),
                          r_re_construct(ast->node2));
    case RE_PIPE:
      return r_re_alter(r_re_construct(ast->node1),
                         r_re_construct(ast->node2));
    case RE_NOT_BOX:
      return r_re_box((RVector*) ast->extra, true);
    case RE_BOX:
      return r_re_box((RVector*) ast->extra, false);
  }
  assert_never();
  return r_re_dot();
}

ReSubExpression r_re_concat(ReSubExpression first, ReSubExpression second)
{
  ReSubExpression result;

  result.input = first.input;
  result.output = second.output;
  /* Now route the outputs of first into input of second */
  r_re_route(first.output, second.input);
  return result;
}

ReSubExpression r_re_dot()
{
  ReSubExpression result;
  NFANode* node = (NFANode*) r_mem_malloc(sizeof(NFANode));
  r_nfanode_init_char(node, 1);
  r_nfanode_add_class(node, R_NFA_CLASS_ANY, 0, 0);
  result.input = node;
  result.output = node;
  return result;
}

ReSubExpression r_re_letter(uni_char letter)
{
  ReSubExpression result;
  NFANode* node = (NFANode*) r_mem_malloc(sizeof(NFANode));
  r_nfanode_init_char(node, 1);
  r_nfanode_add_class(node, R_NFA_CLASS_SINGLE, letter, 0);
  result.input = node;
  result.output = node;
  return result;
}

/* TODO: Like so many other things in this code, this needs proper error
 * handling and reporting. */
bool r_re_match(char* s, NFANode* pattern, bool greedy, int* id, int* length)
{
  char* start = s;
  char* end = s + strlen(s);
  NFAState state;
  r_nfastate_init(&state);
  r_nfastate_activate(&state, pattern);

  /* These are used when greedy = true */
  bool matched = false;
  int matched_len = 0, matched_id = 0;

  while (s != end){
    uni_char letter;
    if (r_utf8_read_uni_char((uint8_t**) (&s), (uint8_t*) end, &letter) != R_SUCCESS)
      return false;

    r_nfastate_step(&state, letter);
    if (not r_nfastate_any(&state)){
      /* Break out of the loop. */
      break;
    }

    /* Go through the active states and find if any are terminating. Match
     * the id with the highest id among the terminating states. This is
     * needed for the operation of the lexer. This means that the rule that
     * was added the latest will get precedence if it matches the same text
     * like an older rule. */
    int i;
    int maximum_id = -1;
    for (i = 0; i < r_nfastate_num(&state); i++){
      if(r_nfastate_active(&state, i)->type == R_NFA_NODE_TERM){
        int this_id = r_nfastate_active(&state, i)->nfa_term.id;
        if (this_id > maximum_id){
          r_debug(R_DETAILED, "%d beats %d", this_id, maximum_id);
          maximum_id = this_id;
          matched_len = s - start;
          matched_id = this_id;
          /* Note the match and continue processing */
          matched = true;
        }

      } /* if state is terminating */
    } /* for state in active states */

    if (matched and not greedy){
      *id = matched_id;
      *length = matched_len;
      return true;
    }
  } /* while (s != end) */

  if (matched){
    *id = matched_id;
    *length = matched_len;
    r_debug(R_DETAILED, "matched %d", matched_id);
    return true;
  }
  return false;
}

RRENode* r_re_new_node(uint8_t type)
{
  RRENode* node = (RRENode*) r_mem_malloc(sizeof(RRENode));
  node->type = type;
  return node;
}

ReSubExpression r_re_plus(ReSubExpression expr)
{
  ReSubExpression result;
  result.input = expr.input;
  NFANode* split_node = (NFANode*) r_mem_malloc(sizeof(NFANode));
  r_nfanode_init_split(split_node, 2);
  result.output = split_node;
  r_nfanode_add_child(split_node, expr.input);
  r_re_route(expr.output, split_node);
  return result;
}

ReSubExpression r_re_question(ReSubExpression expr)
{
  ReSubExpression result;
  NFANode* split_node = (NFANode*) r_mem_malloc(sizeof(NFANode));
  r_nfanode_init_split(split_node, 2);
  r_nfanode_add_child(split_node, expr.input);
  result.input = split_node;
  result.output = split_node;
  split_node->extra = expr.output;
  return result;
}

void r_re_route(NFANode* output, NFANode* target)
{
  while (output != NULL) {
    switch (output->type) {
      case R_NFA_NODE_NOT_CHAR:
      case R_NFA_NODE_CHAR:
        output->nfa_char.next = target;
        break;
      case R_NFA_NODE_SPLIT:
        r_nfanode_add_child(output, target);
        break;
      default: ;
    }
    output = output->extra;
  }; /* while output != NULL */
}

ReSubExpression r_re_star(ReSubExpression expr)
{
  ReSubExpression result;
  NFANode* split_node = (NFANode*) r_mem_malloc(sizeof(NFANode));
  r_nfanode_init_split(split_node, 2);
  r_nfanode_add_child(split_node, expr.input);
  r_re_route(expr.output, split_node);
  result.input = split_node;
  result.output = split_node;
  return result;
}
